from sqlalchemy.orm.session import Session

from airflow.exceptions import AirflowException
from airflow.models.idata_folder import Folder
from airflow.models.idata_task import Task
from airflow.models.idata_workflow import WorkFlow
from airflow.utils.session import provide_session


class TreeNode(object):
    def __init__(self, id, project_id, name, node_type, category, level, children=None, session=None, operator=None):
        self.id = id
        self.project_id = project_id
        self.name = name
        self.node_type = node_type  # 0 folder 1 file
        self.category = category  # 0 task 1 workflow
        self.level = level
        self.children = children
        self.session = session
        self.operator = operator

    def to_json(self):
        return {
            'id': self.id,
            'project_id': self.project_id,
            'name': self.name,
            'node_type': self.node_type,
            'category': self.category,
            'level': self.level,
            'operator': self.operator
        }

    def to_json_tree(self):
        node = self.to_json()
        if self.children:
            node['children'] = [child.to_json_tree() for child in self.children]
        return node

    def set_children(self):
        if not self.node_type == 0:
            return

        folders = Folder.query_children(folder_id=self.id, project_id=self.project_id, session=self.session)
        for folder in folders:
            node = TreeNode(id=folder.id, project_id=folder.project_id, name=folder.name, node_type=0,
                            category=folder.category, level=folder.level, session=self.session)
            node.set_children()
            if self.children is None:
                self.children = [node]
            else:
                self.children.append(node)

        if self.category == 0:
            tasks = Task.query_task(folder_id=self.id, project_id=self.project_id)
            for task in tasks:
                node = TreeNode(id=task.id, project_id=task.project_id, name=task.name, node_type=1,
                                category=0, level=self.level + 1, session=self.session, operator = task.operator)
                if self.children is None:
                    self.children = [node]
                else:
                    self.children.append(node)
        else:
            workflows = WorkFlow.query_workflow(folder_id=self.id, project_id=self.project_id)
            for workflow in workflows:
                node = TreeNode(id=workflow.id, project_id=workflow.project_id, name=workflow.name, node_type=1,
                                category=1, level=self.level + 1, session=self.session)
                if self.children is None:
                    self.children = [node]
                else:
                    self.children.append(node)


@provide_session
def delete(folder_id, session: Session = None):
    sub = Folder.query_children(folder_id=folder_id, session=session)
    tasks = Task.query_task(folder_id=folder_id, session=session)
    wfs = WorkFlow.query_workflow(folder_id=folder_id, session=session)
    if sub or tasks or wfs:
        raise AirflowException('folder is not empty')
    session.query(Folder).filter(Folder.id == folder_id).delete()


@provide_session
def get_tree(project_id, session: Session = None):
    folders = session.query(Folder).filter(Folder.level == 0).all()
    data = []
    for folder in folders:
        node = TreeNode(id=folder.id, project_id=project_id, name=folder.name, node_type=0,
                        category=folder.category, level=folder.level, session=session)
        node.set_children()
        data.append(node.to_json_tree())
    return data
